import time
import tushare as ts
import pandas as pd
import akshare as ak
import datetime
from retrying import retry
from data_resource.data_bases import (Session, engine, engine2, MinuteBars, TsIndicator, StockDaily, StockDailyHFQ,
                                      StockDailyQFQ, IndexDaily)
from data_resource.data_bases import (dtype_basic, dtype_indicators, MacroBondCN, MacroBondUS, MacroCPI, MacroPPI,
                                      MacroSocialFinance)
from sqlalchemy import func
from tqdm import tqdm
from utilities.utilities_func import MyThread, init_ts, isEarningSeason, clear_financials, clear_forecast, \
    bulk_insert_with_orm, write_to_db, get_session
from data_resource.factor_calculate_daily import (factorCalcu_f_return_wgt, f_volatility_n_std,
                                                  f_earningSuprise, f_fama, factorCalcu_f_20dVolSkew, f_rvi_up,
                                                  f_finance_indicator)
from factorPool.golden import golden_daily
from dateutil.relativedelta import relativedelta

pro = init_ts()


# 股票基本信息
def stock_pool_ts():
    _sql = """
    select ticker, short_name, market
    from quant_research.basic_info_stock
    where status = 'L';
    """
    _pool = pd.read_sql(_sql, engine)
    return _pool


# 申万行业分类
def get_sw_industry_classify():
    """获取申万行业分类，1-3级"""
    _r = []
    for i in ['L1', 'L2', 'L3']:
        _a = pro.index_classify(level=i, src='SW2021')
        _r.append(_a)
    _rr = pd.concat(_r)

    return _rr.reset_index(drop=True)


def get_sw_constituent(swlist):
    """获取申万行业成分股"""
    _code = swlist[swlist['level'] == 'L1']['index_code'].unique()
    _r = []
    for _ in _code:
        _n = pro.index_member_all(l1_code=_)
        _r.append(_n)

    _rr = pd.concat(_r)
    return _rr.reset_index(drop=True)


"""
    行情数据接口
"""


@retry(stop_max_attempt_number=3, wait_fixed=3000)
def get_price(code, start_date: str, end_date: str):
    _b = ak.stock_zh_a_hist(symbol=code, period='daily', start_date=start_date,
                            end_date=end_date, adjust='hfq', timeout=30)
    return _b


@retry(stop_max_attempt_number=3, wait_fixed=3000)
def get_1min(code):
    _d = ak.stock_zh_a_hist_min_em(symbol=code, period='1', adjust="hfq")
    return _d


@retry(stop_max_attempt_number=3, wait_fixed=3000)
def ts_bar_daily(code, start_date, end_date, adj='hfq'):
    """ts日行情接口, 支持后复权和不复权行情"""
    assert adj in ['', 'hfq', 'qfq']

    if adj == 'hfq':
        _daily = ts.pro_bar(ts_code=code, asset='E', adj='hfq', freq='D',
                            start_date=start_date, end_date=end_date)
    elif adj == 'qfq':
        _daily = ts.pro_bar(ts_code=code, asset='E', adj='qfq', freq='D',
                            start_date=start_date, end_date=end_date)
    else:
        _daily = ts.pro_bar(ts_code=code, asset='E', freq='D', adj=None,
                            start_date=start_date, end_date=end_date)
    return _daily


@retry(stop_max_attempt_number=3, wait_fixed=3000)
def ts_index_daily(indexcode, start_date, end_date):
    """获取指数日行情，start_date=20250221"""
    _index = pro.index_daily(ts_code=indexcode, start_date=start_date, end_date=end_date)
    return _index


@retry(stop_max_attempt_number=3, wait_fixed=3000)
def ts_indicator_daily(tradedate):
    """获取个股指标日行情，start_date=20250221"""
    _index = pro.daily_basic(ts_code='', trade_date=tradedate)
    return _index


def get_market_1min(stock):
    """
        获取A股1分钟行情数据
        :param stock: 股票池, stock_pool()返回结果
    """
    _session = Session()
    _start = _session.query(func.max(MinuteBars.trading)).scalar()

    if _start is None:
        raise ValueError("数据库中无分钟交易数据，请检查相关连接")
    else:
        _start = _start + datetime.timedelta(days=1)
        _start = _start.replace(hour=0, minute=0, second=0, microsecond=0)

    _end = datetime.datetime.today()

    if _start > _end:
        print(f"数据库中分钟行情数据已是最新数据，无需更新")
        return

    print(f"---------开始获取分钟线数据，行情区间：开始：{_start.date()}, 结束：{_end.date()}-------------")
    results = []
    for i in tqdm(range(len(stock))):

        _code = stock.iloc[i, 0]

        try:
            _data = get_1min(_code)  # 获取后复权数据
        except:
            print(f"获取分钟行情失败:{_code}", flush=True)
            continue

        if _data.empty:
            print(f"{_code}:空分钟行情数据", flush=True)
            continue

        _data['时间'] = pd.to_datetime(_data['时间'])
        # 过滤时间区间内数据
        _data1 = _data.loc[(_data['时间'] > _start) & (_data['时间'] < _end), :].copy()

        if _data1.empty:
            continue

        _data1['code'] = _code
        _data1.drop(columns='均价', inplace=True)

        _data1.rename(columns={
            '时间': 'trade_time', '开盘': 'open', '收盘': 'close', '最高': 'high', '最低': 'low',
            '成交量': 'volumn', '成交额': 'amount'
        })

        results.append(_data1)

    results = pd.concat(results)
    results.reset_index(inplace=True, drop=True)

    results.rename(columns={
        '时间': 'trading', '开盘': 'open', '收盘': 'close', '最高': 'high', '最低': 'low',
        '成交量': 'volumn', '成交额': 'amount'
    }, inplace=True)
    results['trading'] = pd.to_datetime(results['trading'], format='mixed')

    if results.shape[0] < 1300000:
        raise ValueError(f"今日分钟行情数据获取过少，总数据量{results.shape[0]}")
    else:
        print("-------------开始写入分钟线行情数据库，总共有%d条数据-------------" % len(results))
        write_to_db('minute_bars', engine, MinuteBars, results)
        print("-------------分钟线行情数据写入完成---------------------")
    _session.close()

    return


def get_ts_bardaily(fq_type='hfq'):
    """
        根据复权类型获取行情数据
    """
    assert fq_type in ["hfq", '', 'qfq']
    _session = Session()

    if fq_type == 'hfq':
        _start = _session.query(func.max(StockDailyHFQ.trade_date)).scalar()
    elif fq_type == '':
        _start = _session.query(func.max(StockDaily.trade_date)).scalar()
    else:
        _start = _session.query(func.max(StockDailyQFQ.trade_date)).scalar()

    if _start is None:
        raise ValueError("数据库中无日线交易数据，请检查相关连接")
    else:
        _start = _start + datetime.timedelta(days=1)
        _start = _start.strftime("%Y%m%d")

    _end = datetime.datetime.today().date().strftime("%Y%m%d")

    # 获取股票列表
    _sql1 = """
    select ticker, list_date from quant_research.basic_info_stock
    where status='L' and market in ('主板', '创业板', '科创板')
    """
    _pool = pd.read_sql(_sql1, engine)

    result = []
    for num in tqdm(range(len(_pool))):
        _stock = _pool.iloc[num, 0]
        if _stock[0] == '6':
            _code = _stock + '.SH'
        else:
            _code = _stock + '.SZ'
        # _start = _pool.iloc[num, 1].strftime('%Y%m%d')

        # 获取行情数据
        _data = ts_bar_daily(_code, _start, _end, adj=fq_type)

        if _data is None:
            print(f"{_code}无法提取{fq_type}行情数据，时间区间{_start}-{_end}")
        elif _data.empty:
            print(f"{_code}无法提取{fq_type}行情数据，时间区间{_start}-{_end}")
        else:
            result.append(_data)

        if (num % 490 == 0) & (num > 0):
            time.sleep(60)

    result = pd.concat(result)
    result.reset_index(inplace=True, drop=True)
    result.rename(columns={'ts_code': 'ticker'}, inplace=True)
    result['ticker'] = result['ticker'].apply(lambda x: x[0:6])
    result['trade_date'] = pd.to_datetime(result['trade_date'], format='%Y%m%d')

    print("-------------开始写入个股日线行情数据库，总共有%d条数据-------------" % len(result))
    if fq_type == 'hfq':
        write_to_db('market_daily_ts', engine, StockDailyHFQ, result)
    elif fq_type == 'qfq':
        write_to_db('market_daily_qfq', engine, StockDailyQFQ, result)
    else:
        write_to_db('market_daily_none', engine, StockDaily, result)
    print("-------------个股日线行情数据写入完成---------------------")
    _session.close()
    return


def get_ts_indexdaily():
    _session = Session()
    _start = _session.query(func.max(IndexDaily.trade_date)).scalar()

    if _start is None:
        raise ValueError("数据库中无指数日线交易数据，请检查相关连接")
    else:
        _start = _start + datetime.timedelta(days=1)
        _start = _start.strftime("%Y%m%d")
    _end = datetime.datetime.today().date().strftime("%Y%m%d")

    index_name = {
        '000922.CSI': '中证红利',
        '000300.SH': '沪深300',
        '000906.SH': '中证800',
        '000852.SH': '中证1000',
        '000985.CSI': '中证全指',
        '000058.SH': '全指价值',
        '000057.SH': '全指成长',
        '399373.SZ': '大盘价值',
        '399377.SZ': '小盘价值'
    }
    _result = []
    for code in index_name.keys():
        _index = ts_index_daily(code, _start, _end)
        _index['index_name'] = index_name[code]
        if _index is None:
            print(f"指数{code}: 返回空数据，时间区间为{_start}至{_end}")
            continue
        elif _index.empty:
            print(f"指数{code}: 返回空数据，时间区间为{_start}至{_end}")
            continue
        _result.append(_index)

    _result = pd.concat(_result)
    _result.reset_index(drop=True, inplace=True)
    _result['trade_date'] = pd.to_datetime(_result['trade_date']).dt.date

    if _result.empty:
        print(f"所有指数都返回空数据，时间区间为{_start}至{_end}")
    else:
        print("-------------开始写入指数日线行情数据库，总共有%d条数据-------------" % len(_result))
        write_to_db('index_daily', engine, IndexDaily, _result)
        print("-------------指数日线行情数据写入完成---------------------")

    _session.close()
    return


def get_ts_indicator():
    """获取每日个股指标数据"""
    _session = Session()
    try:
        _start = _session.query(func.max(TsIndicator.trade_date)).scalar()

        if _start is None:
            raise ValueError("数据库中无日线交易数据，请检查相关连接")
        else:
            _start = _start + datetime.timedelta(days=1)
        _end = datetime.datetime.today().date()

        if _end > _start:
            # 获取交易日
            _trade = pro.trade_cal(exchange='SSE', start_date=_start.strftime("%Y%m%d"),
                                   end_date=_end.strftime("%Y%m%d"), is_open='1')
            _result = []
            for num, trade in enumerate(_trade['cal_date'].to_list()):
                _data = ts_indicator_daily(trade)
                if _data.empty:
                    print(f"返回空数据，交易日{trade}")
                    continue
                else:
                    _result.append(_data)

                if (num % 190 == 0) & (num > 0):
                    time.sleep(60)
            _result = pd.concat(_result, ignore_index=True)

        else:
            _result = ts_indicator_daily(_end.strftime("%Y%m%d"))

            if _result is None:
                print(f"返回空数据，交易日{_end}")
                return
            elif _result.empty:
                print(f"返回空数据，交易日{_end}")
                return

        _result.rename(columns={'ts_code': 'code'}, inplace=True)
        _result['trade_date'] = pd.to_datetime(_result['trade_date'])
        _result['code'] = _result['code'].str[:6]

        print("-------------开始写入个股指标日线行情数据库，总共有%d条数据-------------" % len(_result))

        with engine.connect() as conn1, engine2.connect() as conn2:
            _result.to_sql(name='indicator_daily', con=conn1, schema='quant_research', index=False,
                           dtype=dtype_indicators, if_exists='append', chunksize=10000, method='multi')

            _result.to_sql(name='indicator_daily', con=conn2, index=False,
                           dtype=dtype_indicators, if_exists='append', chunksize=10000, method='multi')

            print("-------------个股指标日线行情数据写入完成---------------------")
    except Exception as e:
        print(f"发生错误{e}")
        _session.rollback()
        raise e
    finally:
        _session.close()

    return


"""
    ts vip财务数据接口
"""


def get_ts_financials():
    """
        判断当前日期是否为财报季或预测季
    """
    _temp = isEarningSeason(datetime.datetime.today())

    if _temp['forecast']:
        # 写入预测数据
        _forecast = pro.forecast_vip(period=_temp['forecastdate'])
        if _forecast.empty:
            print(f"预测数据为空，时间区间为{_temp['forecastdate']}")
        else:
            _df = clear_forecast(_forecast)
            bulk_insert_with_orm(_df, "financials_forecast", engine,
                                 unique_columns=['ticker', 'ann_date', 'end_date', 'update_flag'])

    if _temp['financials']:
        # 写入报告期财务数据

        _cashflow = []
        _balancesheet = []
        _incomestate = []

        for date in _temp['reportdate']:

            _cashflow1 = pro.cashflow_vip(period=date, report_type=1)
            _cashflow2 = pro.cashflow_vip(period=date, report_type=2)
            _cashflow5 = pro.cashflow_vip(period=date, report_type=5)

            _income1 = pro.income_vip(period=date, report_type=1)
            _income2 = pro.income_vip(period=date, report_type=2)
            _income5 = pro.income_vip(period=date, report_type=5)

            _bs1 = pro.balancesheet_vip(period=date, report_type=1)
            _bs5 = pro.balancesheet_vip(period=date, report_type=5)

            if not _cashflow1.empty:
                _cf1 = clear_financials(_cashflow1)
                bulk_insert_with_orm(_cf1, "financials_cashflow_lt", engine,
                                     unique_columns=['ticker', 'f_ann_date', 'end_date', 'update_flag'])
            if not _cashflow2.empty:
                _cf2 = clear_financials(_cashflow2)
                bulk_insert_with_orm(_cf2, "financials_cashflow_quarter", engine,
                                     unique_columns=['ticker', 'f_ann_date', 'end_date', 'update_flag'])
            if not _cashflow5.empty:
                _cf5 = clear_financials(_cashflow5)
                bulk_insert_with_orm(_cf5, "financials_cashflow_beforeUpdate", engine,
                                     unique_columns=['ticker', 'f_ann_date', 'end_date', 'update_flag'])

            if not _income1.empty:
                _is1 = clear_financials(_income1)
                bulk_insert_with_orm(_is1, "financials_incomeState_lt", engine,
                                     unique_columns=['ticker', 'f_ann_date', 'end_date', 'update_flag'])
            if not _income2.empty:
                _is2 = clear_financials(_income2)
                bulk_insert_with_orm(_is2, "financials_incomeState_quarter", engine,
                                     unique_columns=['ticker', 'f_ann_date', 'end_date', 'update_flag'])
            if not _income5.empty:
                _is5 = clear_financials(_income5)
                bulk_insert_with_orm(_is5, "financials_incomeState_beforeUpdate", engine,
                                     unique_columns=['ticker', 'f_ann_date', 'end_date', 'update_flag'])

            if not _bs1.empty:
                _bs1 = clear_financials(_bs1)
                bulk_insert_with_orm(_bs1, "financials_bs_lt", engine,
                                     unique_columns=['ticker', 'f_ann_date', 'end_date', 'update_flag'])
            if not _bs5.empty:
                _bs5 = clear_financials(_bs5)
                bulk_insert_with_orm(_bs5, "financials_bs_beforeUpdate", engine,
                                     unique_columns=['ticker', 'f_ann_date', 'end_date', 'update_flag'])


"""
    更新个股基本信息
"""


@retry(stop_max_attempt_number=3, wait_fixed=3000)
def get_stock_basic(temp='L'):
    _b = pro.stock_basic(list_status=temp)
    return _b


def update_stock_basic():
    results = []
    temps = ['D', 'P', 'L']
    for temp in temps:
        _data = get_stock_basic(temp)
        if _data.empty:
            continue
        else:
            _data.loc[:, 'status'] = temp
            results.append(_data)
    result = pd.concat(results)
    result1 = result.drop(columns=['ts_code', 'cnspell'])
    result1.rename(columns={'symbol': 'ticker', 'name': 'short_name'}, inplace=True)
    result1.reset_index(inplace=True, drop=True)

    with engine.begin() as connection:
        try:
            result1.to_sql(name='basic_info_stock', con=connection, schema='quant_research', index=False,
                           dtype=dtype_basic,
                           if_exists='replace', chunksize=100000, method='multi')
        except Exception as e:
            print(f"---------- 个股基本信息：basic_info_stock 写入本地数据库失败, {e} ----------")
        else:
            print(
                f"---------- 个股基本信息：basic_info_stock 写入本地数据库成功, 共写入{result1.shape[0]}条数据 ----------")

    with engine2.begin() as connection2:
        try:
            result1.to_sql(
                "basic_info_stock", con=connection2, if_exists='replace', index=False, method="multi",
                chunksize=10000, dtype=dtype_basic
            )
        except Exception as e:
            print(f"---------- 个股基本信息：basic_info_stock 写入远程数据库失败, {e} ----------")
        else:
            print(
                f"---------- 个股基本信息：basic_info_stock 写入远程数据库成功, 共写入{result1.shape[0]}条数据 ----------")

    return


"""
    宏观数据采集
"""


def get_marcro():
    # 获取国债收益率数据
    bond_term_cn = []
    _end = datetime.datetime.today().strftime('%Y%m%d')
    with get_session(engine) as session:
        _start = session.query(func.max(MacroBondCN.trade_date)).scalar()
        assert _start is not None
        _start = _start + datetime.timedelta(days=1)
        _start = _start.strftime('%Y%m%d')

    _bond1 = pro.yc_cb(start_date=_start, end_date=_end, curve_term=10, curve_type="1")
    _bond2 = pro.yc_cb(start_date=_start, end_date=_end, curve_term=5, curve_type="1")
    _bond3 = pro.yc_cb(start_date=_start, end_date=_end, curve_term=1, curve_type="1")
    bond_term_cn.append(_bond1)
    bond_term_cn.append(_bond2)
    bond_term_cn.append(_bond3)
    bond_term_cn = pd.concat(bond_term_cn)
    bond_term_cn.reset_index(inplace=True, drop=True)
    bond_term_cn.rename(columns={'yield': 'yield_rate'}, inplace=True)

    # 获取美债收益率数据
    with get_session(engine) as session:
        _start = session.query(func.max(MacroBondUS.date)).scalar()
        assert _start is not None
        _start = _start + datetime.timedelta(days=1)
        _start = _start.strftime('%Y%m%d')
    _bond_us = pro.us_tycr(start_date=_start, end_date=_end)

    # 获取CPI数据
    _month = datetime.datetime.today()
    with get_session(engine) as session:
        _month_start = session.query(func.max(MacroCPI.month)).scalar()
        assert _month_start is not None
        _month_start = datetime.datetime.strptime(_month_start, '%Y%m')
    if _month_start < _month:
        next_month = _month_start + relativedelta(months=1)
        cpi = pro.cn_cpi(start_m=next_month.strftime('%Y%m'), end_m=_month.strftime('%Y%m'))
    else:
        cpi = pd.DataFrame()

    # 获取PPI数据
    with get_session(engine) as session:
        _month_start = session.query(func.max(MacroPPI.month)).scalar()
        assert _month_start is not None
        _month_start = datetime.datetime.strptime(_month_start, '%Y%m')
    if _month_start < _month:
        next_month = _month_start + relativedelta(months=1)
        ppi = pro.cn_ppi(start_m=next_month.strftime('%Y%m'), end_m=_month.strftime('%Y%m'))
    else:
        ppi = pd.DataFrame()

    # 获取社融数据
    with get_session(engine) as session:
        _month_start = session.query(func.max(MacroSocialFinance.month)).scalar()
        assert _month_start is not None
        _month_start = datetime.datetime.strptime(_month_start, '%Y%m')
    if _month_start < _month:
        next_month = _month_start + relativedelta(months=1)
        sf = pro.sf_month(start_m=next_month.strftime('%Y%m'), end_m=_month.strftime('%Y%m'))
    else:
        sf = pd.DataFrame()

    # 写入数据
    if bond_term_cn.empty:
        print("没有获取到国债收益率数据")
    else:
        write_to_db("bond_term_cn", engine, MacroBondCN, bond_term_cn)

    if _bond_us.empty:
        print("没有获取到美债收益率数据")
    else:
        write_to_db("bond_term_us", engine, MacroBondUS, _bond_us)

    if cpi.empty:
        print("没有获取到CPI数据")
    else:
        write_to_db("macro_cpi_cn", engine, MacroCPI, cpi)

    if ppi.empty:
        print("没有获取到PPI数据")
    else:
        write_to_db("macro_ppi_cn", engine, MacroPPI, ppi)

    if sf.empty:
        print("没有获取到社融数据")
    else:
        write_to_db("macro_social_financing_cn", engine, MacroSocialFinance, sf)


"""
    运行主函数
"""


def run_daily(start, today):
    update_stock_basic()  # 更新个股基本信息
    stock = stock_pool_ts()  # 获取股票池
    get_ts_bardaily(fq_type='hfq')  # 更新日线数据
    get_ts_indexdaily()  # 更新指数日线数据
    get_ts_indicator()  # 更新个股指标日线数据
    get_ts_financials()  # 更新财务指标数据
    get_marcro()  # 更新宏观数据

    # 定义任务列表
    tasks = [
        # ("MarketData", get_market_1min, (stock,)),  # 采集分钟行情数据
        ("ReturnWgtFactor", factorCalcu_f_return_wgt, ()),  # 计算换手率加权区间涨跌幅因子
        ("VolatilityFactor", f_volatility_n_std, ()),  # 计算区间波动率因子
        ("EarningSurpriseFactor", f_earningSuprise, ()),  # 计算营收超预期因子
        ("FamaFactor", f_fama, ()),  # 计算fama因子
        ("VolSkewFactor", factorCalcu_f_20dVolSkew, ()),  # 计算成交量加权收益率偏度因子
        ("RVIFactor", f_rvi_up, ()),  # 计算RVI波动率因子
        ("FinancialIndicatgor", f_finance_indicator, ())  # 更新财务指标
        # ("Golden", golden_daily, (start, today))  # 黄金台计算
    ]

    # 创建线程对象
    threads = []
    for name, func, args in tasks:
        thread = MyThread(func, args=args, name=name)
        threads.append((name, thread))

    # 启动所有线程
    for name, thread in threads:
        print(f"Starting task: {name}")
        thread.start()

    # 等待所有线程完成并收集结果
    errors = []
    for name, thread in threads:
        thread.join()
        if hasattr(thread, 'exception') and thread.exception:
            error_info = f"Task '{name}' failed with error: {thread.exception}"
            errors.append(error_info)
            print(error_info)
        else:
            print(f"Task '{name}' completed successfully")

    # 如果有任何错误，可以选择抛出异常
    if errors:
        print("Some tasks failed:")
        for error in errors:
            print(error)

    print("------------ 开始黄金台信号计算 --------")
    golden_daily(start, today)


def run_weekend():
    get_ts_bardaily(fq_type='')  # 更新无复权日线数据
    get_ts_bardaily(fq_type='qfq')  # 更新前复权日线数据


if __name__ == '__main__':
    get_ts_financials()
